const DOCKER_HUB = 'https://registry-1.docker.io';
const PROXY_MAP: Record<string, string> = {
	[`docker.`]: DOCKER_HUB,
	[`quay.`]: 'https://quay.io',
	[`gcr.`]: 'https://gcr.io',
	[`k8s-gcr.`]: 'https://k8s.gcr.io',
	[`k8s.`]: 'https://registry.k8s.io',
	[`ghcr.`]: 'https://ghcr.io',
	[`cloudsmith.`]: 'https://docker.cloudsmith.io',
	[`ecr.`]: 'https://public.ecr.aws'
};

function getUpstream(hostname: string): string {
	for (const [prefix, upstream] of Object.entries(PROXY_MAP)) {
		if (hostname.startsWith(prefix)) {
			return upstream;
		}
	}

	throw new Error(`No upstream found for ${hostname}`);
}

async function upstream_v2(url: URL, authorization: string | null): Promise<Response> {
	const newUrl = URL.parse('/v2/', getUpstream(url.hostname));
	if (!newUrl) {
		return new Response('Invalid upstream', { status: 500 });
	}

	const headers = new Headers();
	if (authorization) {
		headers.set('Authorization', authorization);
	}

	const res = await fetch(newUrl.toString(), { method: 'GET', headers, redirect: 'follow' });
	if (res.status === 401) {
		return new Response(res.statusText, {
			status: res.status,
			headers: { 'WWW-Authenticate': `Bearer realm="https://${url.hostname}/v2/auth",service="cloudflare-docker-proxy` }
		});
	}

	return res;
}

async function upstream_v2_auth(url: URL, authorization: string | null): Promise<Response> {
	const upstream = getUpstream(url.hostname);
	const newUrl = URL.parse('/v2/', upstream);
	if (!newUrl) {
		return new Response('Invalid upstream', { status: 500 });
	}

	const res = await fetch(newUrl.toString(), { method: 'GET', redirect: 'follow' });
	if (res.status !== 401) {
		return res;
	}

	const wwwAuthenticate = res.headers.get('WWW-Authenticate');
	if (!wwwAuthenticate) {
		return res;
	}

	const authenticate = parseAuthenticate(wwwAuthenticate);
	const isDockerHub = upstream === DOCKER_HUB;
	let scope = url.searchParams.get('scope');

	if (!!scope && isDockerHub) {
		let scopeParts = scope.split(':');
		if (scopeParts.length === 3 && !scopeParts[1].includes('/')) {
			scopeParts[1] = 'library/' + scopeParts[1];
			scope = scopeParts.join(':');
		}
	}

	const realmUrl = new URL(authenticate.realm);

	if (authenticate.service.length) {
		realmUrl.searchParams.set('service', authenticate.service);
	}

	if (scope) {
		realmUrl.searchParams.set('scope', scope);
	}

	const headers = new Headers();
	if (authorization) {
		headers.set('Authorization', authorization);
	}

	return await fetch(realmUrl, { method: 'GET', headers });
}

function parseAuthenticate(authenticate: string) {
	// sample: Bearer realm="https://auth.ipv6.docker.com/token",service="registry.docker.io"
	// match strings after =" and before "
	const reg = /(?<==")(?:\\.|[^"\\])*(?=")/g;
	const matches = authenticate.match(reg);
	if (matches == null || matches.length < 2) {
		throw new Error(`invalid Www-Authenticate Header: ${authenticate}`);
	}
	return {
		realm: matches[0],
		service: matches[1]
	};
}

export default {
	async fetch(request, env, ctx): Promise<Response> {
		try {

			const url = new URL(request.url);
			const authorization = request.headers.get('Authorization');

			if (url.pathname === '/v2/') {
				return await upstream_v2(url, authorization);
			}

			if (url.pathname === '/v2/auth') {
				return await upstream_v2_auth(url, authorization);
			}


			const upstream = getUpstream(url.hostname);
			const isDockerHub = upstream === DOCKER_HUB;

			if (isDockerHub) {
				const pathParts = url.pathname.split('/');
				if (pathParts.length === 5) {
					pathParts.splice(2, 0, 'library');
					const redirectUrl = new URL(url);
					redirectUrl.pathname = pathParts.join('/');
					return Response.redirect(redirectUrl.toString(), 302);
				}
			}

			const proxyUrl = URL.parse(url.pathname, upstream);
			if (!proxyUrl) {
				return new Response('Invalid upstream', { status: 500 });
			}

			const req = new Request(proxyUrl, {
				method: request.method,
				headers: request.headers,
				redirect: 'follow'
			});

			const res = await fetch(req);
			if (res.status === 401) {
				return new Response(res.statusText, {
					status: res.status,
					headers: { 'WWW-Authenticate': `Bearer realm="https://${url.hostname}/v2/auth",service="cloudflare-docker-proxy` }
				});
			}

			return res;
		} catch (e: any) {
			return new Response(e.message, { status: 500 });
		}
	}
} satisfies ExportedHandler<Env>;
